use slotmap::{new_key_type, SlotMap};
use super::{function, basicblock, instr};

new_key_type!{
    pub struct InstrRef;
    pub struct BasicBlockRef;
    pub struct FunctionRef;
    // ModuleRef,
}

#[derive(Debug, Default)]
pub struct MIRSlotMap {
    pub funcs : SlotMap<FunctionRef, function::func>,
    pub bbs : SlotMap<BasicBlockRef, basicblock::bb>,
    pub instrs : SlotMap<InstrRef, instr::MirInst>
}

impl MIRSlotMap {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_func(&mut self, f: function::func) -> FunctionRef {
        self.funcs.insert(f)
    }

    pub fn add_bb(&mut self, b: basicblock::bb) -> BasicBlockRef {
        self.bbs.insert(b)
    }

    pub fn add_instr(&mut self, i: instr::MirInst) -> InstrRef {
        self.instrs.insert(i)
    }

    pub fn get_func(&self, r: FunctionRef) -> Option<&function::func> {
        self.funcs.get(r)
    }

    pub fn get_bb(&self, r: BasicBlockRef) -> Option<&basicblock::bb> {
        self.bbs.get(r)
    }

    pub fn get_instr(&self, r: InstrRef) -> Option<&instr::MirInst> {
        self.instrs.get(r)
    }

    pub fn get_mut_func(&mut self, r: FunctionRef) -> Option<&mut function::func> {
        self.funcs.get_mut(r)
    }

    pub fn get_mut_bb(&mut self, r: BasicBlockRef) -> Option<&mut basicblock::bb> {
        self.bbs.get_mut(r)
    }

    pub fn get_mut_instr(&mut self, r: InstrRef) -> Option<&mut instr::MirInst> {
        self.instrs.get_mut(r)
    }
}
